#include <cstdlib>
#include <cstring>
#include <errno.h>
#include <pthread.h>
#include <stdio.h>
#include <sys/poll.h>
#include <unistd.h>

#include "rk_comm_rgn.h"
#include "rk_comm_vi.h"
#include "rk_comm_vo.h"
#include "rk_common.h"
#include "rk_debug.h"
#include "rk_defines.h"
#include "rk_mpi_cal.h"
#include "rk_mpi_mb.h"
#include "rk_mpi_mmz.h"
#include "rk_mpi_rgn.h"
#include "rk_mpi_sys.h"
#include "rk_mpi_venc.h"
#include "rk_mpi_vi.h"
#include "rk_mpi_vo.h"
#include "rk_mpi_vpss.h"
#include "test_comm_argparse.h"
#include "test_comm_utils.h"
#include "test_common.h"
#include "udp_sender.hpp"
#include <arpa/inet.h>

#define TEST_VENC_MAX 2
#define TEST_WITH_FD 0
#define TEST_WITH_FD_SWITCH 0

#define H264_VENC_CHN 0
#define JPEG_VENC_CHN 1

UdpSender udpSender;

typedef struct _rkTestVencCfg {
    RK_BOOL bOutDebugCfg;
    VENC_CHN_ATTR_S stAttr;
    RK_CHAR dstFilePath[128];
    RK_CHAR dstFileName[128];
    RK_S32 s32ChnId;
    FILE *fp;
    RK_S32 selectFd;
} TEST_VENC_CFG;

typedef struct rkVPSS_CFG_S {
    const char *dstFilePath;
    RK_S32 s32DevId;
    RK_S32 s32ChnId;
    RK_U32 u32VpssChnCnt;
    VPSS_GRP_ATTR_S stGrpVpssAttr;
    VPSS_CHN_ATTR_S stVpssChnAttr[VPSS_MAX_CHN_NUM];
} VPSS_CFG_S;

typedef enum rkTestVIMODE_E {
    TEST_VI_MODE_VI_ONLY = 0,
    TEST_VI_MODE_BIND_VENC = 1,
    TEST_VI_MODE_BIND_VENC_MULTI = 2,
    TEST_VI_MODE_BIND_VPSS_BIND_VENC = 3,
    TEST_VI_MODE_BIND_VO = 4,
    TEST_VI_MODE_MUTI_VI = 5,
} TEST_VI_MODE_E;

typedef struct rkRGN_CFG_S {
    RGN_ATTR_S stRgnAttr;
    RGN_CHN_ATTR_S stRgnChnAttr;
} RGN_CFG_S;

typedef struct _rkMpiVICtx {
    RK_S32 width;
    RK_S32 height;
    RK_S32 devId;
    RK_S32 pipeId;
    RK_S32 channelId;
    RK_S32 loopCountSet;
    RK_S32 selectFd;
    RK_BOOL bFreeze;
    RK_BOOL bEnRgn;
    RK_S32 s32RgnCnt;
    RK_S32 rgnType;
    RK_BOOL bUserPicEnabled;
    RK_BOOL bGetConnecInfo;
    RK_BOOL bGetEdid;
    RK_BOOL bSetEdid;
    COMPRESS_MODE_E enCompressMode;
    VI_DEV_ATTR_S stDevAttr;
    VI_DEV_BIND_PIPE_S stBindPipe;
    VI_CHN_ATTR_S stChnAttr;
    VI_SAVE_FILE_INFO_S stDebugFile;
    VIDEO_FRAME_INFO_S stViFrame;
    VI_CHN_STATUS_S stChnStatus;
    VI_USERPIC_ATTR_S stUsrPic;
    TEST_VI_MODE_E enMode;
    const char *aEntityName;
    // for vi
    RGN_CFG_S stViRgn;
    // for venc
    TEST_VENC_CFG stVencCfg[TEST_VENC_MAX];
    VENC_STREAM_S stFrame[TEST_VENC_MAX];
    VPSS_CFG_S stVpssCfg;
    // for vo
    VO_LAYER s32VoLayer;
    VO_DEV s32VoDev;
} TEST_VI_CTX_S;

static RK_S32 test_vi_init(TEST_VI_CTX_S *ctx) {
    RK_S32 s32Ret = RK_FAILURE;

    // 0. get dev config status  默认通道0 绑定到 dev/viedo62
    s32Ret = RK_MPI_VI_GetDevAttr(ctx->devId, &ctx->stDevAttr);
    if (s32Ret == RK_ERR_VI_NOT_CONFIG) {
        // 0-1.config dev
        s32Ret = RK_MPI_VI_SetDevAttr(ctx->devId, &ctx->stDevAttr);
        if (s32Ret != RK_SUCCESS) {
            RK_LOGE("RK_MPI_VI_SetDevAttr %x", s32Ret);
            goto __FAILED;
        }
    } else {
        RK_LOGE("RK_MPI_VI_SetDevAttr already");
    }
    // 1.get  dev enable status
    s32Ret = RK_MPI_VI_GetDevIsEnable(ctx->devId);
    if (s32Ret != RK_SUCCESS) {
        // 1-2.enable dev
        s32Ret = RK_MPI_VI_EnableDev(ctx->devId);
        if (s32Ret != RK_SUCCESS) {
            RK_LOGE("RK_MPI_VI_EnableDev %x", s32Ret);
            goto __FAILED;
        }
        // 1-3.bind dev/pipe
        ctx->stBindPipe.u32Num = ctx->pipeId;
        ctx->stBindPipe.PipeId[0] = ctx->pipeId;
        s32Ret = RK_MPI_VI_SetDevBindPipe(ctx->devId, &ctx->stBindPipe);
        if (s32Ret != RK_SUCCESS) {
            RK_LOGE("RK_MPI_VI_SetDevBindPipe %x", s32Ret);
            goto __FAILED;
        }
    } else {
        RK_LOGE("RK_MPI_VI_EnableDev already");
    }
    // 2.config channel
    ctx->stChnAttr.stSize.u32Width = ctx->width;
    ctx->stChnAttr.stSize.u32Height = ctx->height;
    ctx->stChnAttr.enCompressMode = ctx->enCompressMode;
    s32Ret = RK_MPI_VI_SetChnAttr(ctx->pipeId, ctx->channelId, &ctx->stChnAttr);
    if (s32Ret != RK_SUCCESS) {
        RK_LOGE("RK_MPI_VI_SetChnAttr %x", s32Ret);
        goto __FAILED;
    }

    /* test user picture */

    // ctx->stViRgn.stRgnAttr.enType = (RGN_TYPE_E)ctx->rgnType;
    // ctx->stViRgn.stRgnChnAttr.bShow = RK_TRUE;

    // open fd before enable chn will be better

    // 3.enable channel
    s32Ret = RK_MPI_VI_EnableChn(ctx->pipeId, ctx->channelId);
    if (s32Ret != RK_SUCCESS) {
        RK_LOGE("RK_MPI_VI_EnableChn %x", s32Ret);
        goto __FAILED;
    }

    // 4.save debug file
    if (ctx->stDebugFile.bCfg) {
        s32Ret = RK_MPI_VI_ChnSaveFile(ctx->pipeId, ctx->channelId, &ctx->stDebugFile);
        RK_LOGD("RK_MPI_VI_ChnSaveFile %x", s32Ret);
    }

__FAILED:
    return s32Ret;
}
void init_venc_cfg(TEST_VI_CTX_S *ctx, RK_U32 u32Ch, RK_CODEC_ID_E enType) {
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.enType = enType; ////一路264 还缺一路
    ctx->stVencCfg[u32Ch].s32ChnId = u32Ch;
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.enPixelFormat = ctx->stChnAttr.enPixelFormat;
    ctx->stVencCfg[u32Ch].stAttr.stRcAttr.enRcMode = VENC_RC_MODE_H264AVBR;
    // ctx->stVencCfg[u32Ch].stAttr.stRcAttr.stH264Cbr.u32Gop = 2;
    ctx->stVencCfg[u32Ch].stAttr.stRcAttr.stH264Vbr.u32Gop = 30;
    ctx->stVencCfg[u32Ch].stAttr.stRcAttr.stH264Vbr.u32SrcFrameRateNum = 30;
    ctx->stVencCfg[u32Ch].stAttr.stRcAttr.stH264Vbr.u32SrcFrameRateDen = 1;
    ctx->stVencCfg[u32Ch].stAttr.stRcAttr.stH264Vbr.fr32DstFrameRateNum = 30;
    ctx->stVencCfg[u32Ch].stAttr.stRcAttr.stH264Vbr.fr32DstFrameRateDen = 1;
    ctx->stVencCfg[u32Ch].stAttr.stRcAttr.stH264Vbr.u32BitRate = 4000;
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.u32PicWidth = ctx->width;
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.u32PicHeight = ctx->height;
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.u32VirWidth = ctx->width;
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.u32VirHeight = ctx->height;
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.u32StreamBufCnt = 5;
    ctx->stVencCfg[u32Ch].stAttr.stVencAttr.u32BufSize = ctx->width * ctx->height * 3 / 2;
}
static RK_S32 create_venc(TEST_VI_CTX_S *ctx, RK_U32 u32Ch) {
    VENC_RECV_PIC_PARAM_S stRecvParam;

    stRecvParam.s32RecvPicNum = ctx->loopCountSet;
    RK_MPI_VENC_CreateChn(u32Ch, &ctx->stVencCfg[u32Ch].stAttr);
    RK_MPI_VENC_StartRecvFrame(u32Ch, &stRecvParam);

    return RK_SUCCESS;
}
static RK_S32 test_vi_bind_venc_loop(TEST_VI_CTX_S *ctx) {
    MPP_CHN_S stSrcChn, stDestChn[TEST_VENC_MAX];
    RK_S32 loopCount = 0;
    void *pData = RK_NULL;
    RK_S32 s32Ret = RK_FAILURE;
    RK_U32 i;
    RK_U32 u32DstCount = ((ctx->enMode == TEST_VI_MODE_BIND_VENC_MULTI) ? 2 : 1);

    s32Ret = test_vi_init(ctx);
    if (s32Ret != RK_SUCCESS) {
        RK_LOGE("vi %d:%d init failed:%x", ctx->devId, ctx->channelId, s32Ret);
        goto __FAILED;
    }

    /* venc */
    // 绑定多路 vnec h264 jpeg
    for (i = 0; i < u32DstCount; i++) {
        // venc  init and create
        if (i == 0) {
            init_venc_cfg(ctx, i, RK_VIDEO_ID_AVC);
        } else {
            init_venc_cfg(ctx, i, RK_VIDEO_ID_JPEG);
        }

        s32Ret = create_venc(ctx, ctx->stVencCfg[i].s32ChnId);
        if (s32Ret != RK_SUCCESS) {
            RK_LOGE("create %d ch venc failed", ctx->stVencCfg[i].s32ChnId);
            return s32Ret;
        }
        // bind vi to venc
        stSrcChn.enModId = RK_ID_VI;
        stSrcChn.s32DevId = ctx->devId;
        stSrcChn.s32ChnId = ctx->channelId;

        stDestChn[i].enModId = RK_ID_VENC;
        stDestChn[i].s32DevId = i;
        stDestChn[i].s32ChnId = ctx->stVencCfg[i].s32ChnId;

        s32Ret = RK_MPI_SYS_Bind(&stSrcChn, &stDestChn[i]);
        if (s32Ret != RK_SUCCESS) {
            RK_LOGE("create %d ch venc failed", ctx->stVencCfg[i].s32ChnId);
            goto __FAILED;
        }
        ctx->stFrame[i].pstPack = reinterpret_cast<VENC_PACK_S *>(malloc(sizeof(VENC_PACK_S)));
#if TEST_WITH_FD
        ctx->stVencCfg[i].selectFd = RK_MPI_VENC_GetFd(ctx->stVencCfg[i].s32ChnId);
        RK_LOGE("venc chn:%d, ctx->selectFd:%d ", ctx->stVencCfg[i].s32ChnId, ctx->stVencCfg[i].selectFd);
#endif
    }

    while (loopCount < ctx->loopCountSet) {
        for (i = 0; i < u32DstCount; i++) {
#if TEST_WITH_FD
            test_vi_poll_event(-1, ctx->stVencCfg[i].selectFd);
#endif
            // freeze test
            RK_MPI_VI_SetChnFreeze(ctx->pipeId, ctx->channelId, ctx->bFreeze);

            // 对于JPEG通道，每次捕获都创建一个新文件
            if (ctx->stVencCfg[i].stAttr.stVencAttr.enType == RK_VIDEO_ID_JPEG) {
                char jpegFilename[256];
                snprintf(jpegFilename, sizeof(jpegFilename), "/%s/venc_%d_%d.jpeg",
                         ctx->stVencCfg[i].dstFilePath, i, loopCount); // 使用loopCount确保文件名唯一
                ctx->stVencCfg[i].fp = fopen(jpegFilename, "wb");
                if (ctx->stVencCfg[i].fp == NULL) {
                    RK_LOGE("Cannot open file %s for writing JPEG\n", jpegFilename);
                    continue; // 跳过这次循环
                }
            }

            s32Ret = RK_MPI_VENC_GetStream(ctx->stVencCfg[i].s32ChnId, &ctx->stFrame[i], -1);
            if (s32Ret == RK_SUCCESS) {

                if (ctx->stVencCfg[i].bOutDebugCfg) {
                    pData = RK_MPI_MB_Handle2VirAddr(ctx->stFrame[i].pstPack->pMbBlk);
                    if (ctx->stVencCfg[i].stAttr.stVencAttr.enType == RK_VIDEO_ID_AVC) {
                        // 创建一个 sockaddr_in 结构体,表示客户端的地址信息
                        sockaddr_in clientAddr;
                        memset(&clientAddr, 0, sizeof(clientAddr));
                        clientAddr.sin_family = AF_INET;
                        clientAddr.sin_port = htons(4569);                        // 替换为实际的客户端端口号
                        inet_pton(AF_INET, "172.16.31.46", &clientAddr.sin_addr); // 替换为实际的客户端IP地址

                        // 通过 UDP 发送编码后的数据
                        if (!udpSender.sendData(pData, ctx->stFrame[i].pstPack->u32Len, clientAddr)) {
                            RK_LOGE("Failed to send data via UDP");
                        }
                    }
                    fwrite(pData, 1, ctx->stFrame[i].pstPack->u32Len, ctx->stVencCfg[i].fp);
                    fflush(ctx->stVencCfg[i].fp);
                }
                RK_U64 nowUs = TEST_COMM_GetNowUs();

                RK_LOGD("chn:%d, loopCount:%d enc->seq:%d wd:%d pts=%lld delay=%lldus\n", i, loopCount,
                        ctx->stFrame[i].u32Seq, ctx->stFrame[i].pstPack->u32Len,
                        ctx->stFrame[i].pstPack->u64PTS / 1000,
                        nowUs - ctx->stFrame[i].pstPack->u64PTS);
                s32Ret = RK_MPI_VENC_ReleaseStream(ctx->stVencCfg[i].s32ChnId, &ctx->stFrame[i]);
                if (s32Ret != RK_SUCCESS) {
                    RK_LOGE("RK_MPI_VENC_ReleaseStream fail %x", s32Ret);
                }
                loopCount++;
            } else {
                RK_LOGE("RK_MPI_VI_GetChnFrame fail %x", s32Ret);
            }
        }
        usleep(10 * 1000);
    }

__FAILED:
    for (i = 0; i < u32DstCount; i++) {
        s32Ret = RK_MPI_SYS_UnBind(&stSrcChn, &stDestChn[i]);
        if (s32Ret != RK_SUCCESS) {
            RK_LOGE("RK_MPI_SYS_UnBind fail %x", s32Ret);
        }
    }
    // 5. disable one chn
    s32Ret = RK_MPI_VI_DisableChn(ctx->pipeId, ctx->channelId);
    RK_LOGE("RK_MPI_VI_DisableChn %x", s32Ret);

    for (i = 0; i < u32DstCount; i++) {
        s32Ret = RK_MPI_VENC_StopRecvFrame(ctx->stVencCfg[i].s32ChnId);
        if (s32Ret != RK_SUCCESS) {
            return s32Ret;
        }
        RK_LOGE("destroy enc chn:%d", ctx->stVencCfg[i].s32ChnId);
        s32Ret = RK_MPI_VENC_DestroyChn(ctx->stVencCfg[i].s32ChnId);
        if (s32Ret != RK_SUCCESS) {
            RK_LOGE("RK_MPI_VDEC_DestroyChn fail %x", s32Ret);
        }
    }

    // 6.disable dev(will diabled all chn)
    s32Ret = RK_MPI_VI_DisableDev(ctx->devId);
    RK_LOGE("RK_MPI_VI_DisableDev %x", s32Ret);
    for (i = 0; i < u32DstCount; i++) {
        if (ctx->stFrame[i].pstPack)
            free(ctx->stFrame[i].pstPack);
        if (ctx->stVencCfg[i].fp)
            fclose(ctx->stVencCfg[i].fp);
    }
    return s32Ret;
}

int main() {
    RK_S32 s32Ret = RK_SUCCESS;
    VI_CHN viChn = 0;
    VENC_CHN h264VencChn = H264_VENC_CHN;
    VENC_CHN jpegVencChn = JPEG_VENC_CHN;

    // 在初始化部分
    if (!udpSender.initialize(4569)) {
        RK_LOGE("Failed to initialize UDP sender");
        return RK_FAILURE;
    } else {
        RK_LOGE("sucess  to initialize UDP sender 4567");
    }

    TEST_VI_CTX_S *ctx;
    ctx = reinterpret_cast<TEST_VI_CTX_S *>(malloc(sizeof(TEST_VI_CTX_S)));
    memset(ctx, 0, sizeof(TEST_VI_CTX_S));
    ctx->width = 1920;
    ctx->height = 1080;
    ctx->devId = 3;
    ctx->pipeId = ctx->devId;
    ctx->channelId = 3;
    ctx->loopCountSet = 10000;

    // TEST_VI_MODE_BIND_VENC_MULTI
    //  多通道
    ctx->enMode = TEST_VI_MODE_BIND_VENC_MULTI; // TEST_VI_MODE_BIND_VENC
    ctx->stChnAttr.stIspOpt.u32BufCount = 3;
    ctx->stChnAttr.stIspOpt.enMemoryType = VI_V4L2_MEMORY_TYPE_DMABUF;
    ctx->stChnAttr.stIspOpt.enCaptureType = VI_V4L2_CAPTURE_TYPE_VIDEO_CAPTURE;
    ctx->stChnAttr.u32Depth = 2;
    ctx->aEntityName = "/dev/video53"; // dev/video44
    ctx->stChnAttr.enPixelFormat = RK_FMT_YUV420SP;
    ctx->stChnAttr.stFrameRate.s32SrcFrameRate = -1;
    ctx->stChnAttr.stFrameRate.s32DstFrameRate = -1;
    ctx->bEnRgn = RK_FALSE;
    ctx->s32RgnCnt = 1;
    ctx->rgnType = RGN_BUTT;
    RK_LOGE("test running enter!");

    // ctx->enCompressMode = COMPRESS_AFBC_16x16;//帧压缩

    // 启动debug文件
    ctx->stDebugFile.bCfg = RK_TRUE;

    if (ctx->stDebugFile.bCfg) {
        if (ctx->enMode == TEST_VI_MODE_BIND_VENC || ctx->enMode == TEST_VI_MODE_BIND_VPSS_BIND_VENC) {
            ctx->stVencCfg[0].bOutDebugCfg = ctx->stDebugFile.bCfg;
        } else if (ctx->enMode == TEST_VI_MODE_BIND_VENC_MULTI) {
            ctx->stVencCfg[0].bOutDebugCfg = ctx->stDebugFile.bCfg;
            ctx->stVencCfg[1].bOutDebugCfg = ctx->stDebugFile.bCfg;
        }
        memcpy(&ctx->stDebugFile.aFilePath, "/userdata/hx/test_data", strlen("/userdata/hx/test_data"));
        snprintf(ctx->stDebugFile.aFileName, MAX_VI_FILE_PATH_LEN,
                 "test_%d_%d_%d.bin", ctx->devId, ctx->pipeId, ctx->channelId);
    }
    // 配置测试文件
    for (int i = 0; i < ctx->enMode; i++) {
        if (ctx->stVencCfg[i].bOutDebugCfg) {
            char name[256] = {0};
            if (i == 0) { // h264
                memcpy(&ctx->stVencCfg[i].dstFilePath, "userdata/hx/test_data3/h264", strlen("userdata/hx/test_data3/h264"));
                snprintf(ctx->stVencCfg[i].dstFileName, sizeof(ctx->stVencCfg[i].dstFileName),
                         "venc_%d.bin", i);
                snprintf(name, sizeof(name), "/%s/%s",
                         ctx->stVencCfg[i].dstFilePath, ctx->stVencCfg[i].dstFileName);
                ctx->stVencCfg[i].fp = fopen(name, "wb");
                if (ctx->stVencCfg[i].fp == RK_NULL) {
                    RK_LOGE("chn %d can't open file %s in get picture thread!\n", i, name);
                    goto __FAILED;
                }
            } else { // jpeg
                memcpy(&ctx->stVencCfg[i].dstFilePath, "userdata/hx/test_data3/jpeg", strlen("userdata/hx/test_data3/jpeg"));
                snprintf(ctx->stVencCfg[i].dstFileName, sizeof(ctx->stVencCfg[i].dstFileName),
                         "venc_%d.jpeg", i);
                snprintf(name, sizeof(name), "/%s/%s",
                         ctx->stVencCfg[i].dstFilePath, ctx->stVencCfg[i].dstFileName);
                ctx->stVencCfg[i].fp = fopen(name, "wb");
                if (ctx->stVencCfg[i].fp == RK_NULL) {
                    RK_LOGE("chn %d can't open file %s in get picture thread!\n", i, name);
                    goto __FAILED;
                }
            }
        }
    }

    RK_LOGE("test running enter ctx->aEntityName=%s!", ctx->aEntityName);
    if (ctx->aEntityName != RK_NULL)
        memcpy(ctx->stChnAttr.stIspOpt.aEntityName, ctx->aEntityName, strlen(ctx->aEntityName));

    // 1. 初始化系统

    if (RK_MPI_SYS_Init() != RK_SUCCESS) {
        RK_LOGE("rk mpi sys init fail!");
        goto __FAILED;
    }
    switch (ctx->enMode) {
    case TEST_VI_MODE_BIND_VENC:
    case TEST_VI_MODE_BIND_VENC_MULTI:
        s32Ret = test_vi_bind_venc_loop(ctx);
        break;
    default:
        RK_LOGE("no support such test mode:%d", ctx->enMode);
        break;
    }

__FAILED:
    RK_LOGE("test running exit:%d", s32Ret);
    RK_MPI_SYS_Exit();
__FAILED2:
    if (ctx) {
        free(ctx);
        ctx = RK_NULL;
    }
    return 0;
}